"""
IP端口扫描
IP port scan
"""

from common import *


class PortScan:
    def __init__(self, ip, semaphore, port=None, is_all=False,
                 save_path=None):
        """
        定义相关参数  Define related parameters
        :param ip: 扫描的目标IP Scan target ip
        :param port:    自定义端口 Custom ports
        :param is_all: 是否扫描所有端口 Whether scan all tcp ports
        :param semaphore: 并发数   Number of concurrent
        :param save_path: 保存文件的路径   Save file path
        """

        self._save_path = save_path
        self._semaphore = semaphore
        self.target_ip = ip
        # 默认web常用端口 Default common web ports.
        self.web_ports = [21, 22, 23, 25, 69, 80, 81, 89, 90, 110, 443,
                          1080, 1158, 1433, 1434, 1521, 2100, 3306, 3389,
                          7001, 8080, 9080, 9090]
        # 所有端口 All ports.
        self.all_ports = range(65535) if is_all else None
        # 自定义端口 Custom ports.
        self.custom_ports = port
        # 创建输出表格    Create output form.
        self.table = PrettyTable(['port', 'state'])

    async def _port_scan(self, port):
        """
        端口扫描    Port scan
        :param port:    自定义端口 Custom ports
        :return:
        """
        # sem = self._semaphore
        # print(sem)
        async with self._semaphore:
            try:
                if await asyncio.wait_for(asyncio.open_connection(self.target_ip, port), timeout=0.2):
                    self.table.add_row([port, 'open'])
            except:
                pass

    def scan_options(self):
        """
        扫描选项    Scan options
        :return:
        """
        # 如果没有自定义端口或者扫描全部端口就默认扫描web端口
        # Default scan web port if not have custom ports or not scan all Tcp ports
        if not self.custom_ports:
            if not self.all_ports:
                batch_run(func=self._port_scan, parma=self.web_ports,
                          semaphore=self._semaphore, table=self.table,
                          save_path=self._save_path, save_func=self.save_result)
            else:
                batch_run(func=self._port_scan, parma=self.all_ports,
                          semaphore=self._semaphore, table=self.table,
                          save_path=self._save_path, save_func=self.save_result)
        else:
            batch_run(func=self._port_scan, parma=self.custom_ports,
                      semaphore=self._semaphore, table=self.table,
                      save_path=self._save_path, save_func=self.save_result)

    def save_result(self):
        """
        保存扫描结果  Save result data.
        :return:
        """
        if self._save_path:
            with open(self._save_path+'.csv', 'w', newline='') as fp:
                # 将表格转换为json字符串用json.loads读取出来
                # Load a table into a json string with json.loads.
                table = json.loads(self.table.get_csv_string())
                # Take out the list head.
                # 取出表头
                header = table[0]
                table.remove(table[0])
                # 写入数据
                # Write data.
                writer = csv.DictWriter(fp, header)
                writer.writeheader()
                writer.writerows(table)
            success_color.print('保存成功\t', os.getcwd()+'\\' + self._save_path + '.csv')
